'''
Created on Oct 12, 2009

@author: mkiyer
'''

from optparse import OptionParser
from veggie.app.rnaseq.tophat import TopHat
from veggie.sample.sampledb2 import get_sampledb
from veggie.sample.samplegroup import parse_samplegroups_xml
import sys
import logging
import os
import glob
import re

def find_pe_sequence_files(flowcell, lane, sequence_path,
                           prefix_fmt="%s_%s",
                           seq_suffix="best_sequence.txt",
                           bowtie_quals_suffix='best_bowtiequals.txt'):
    analysis_id = prefix_fmt % (flowcell, lane)
    # build a string to use to search for sequence files    
    mate1_search_string = ''.join([analysis_id, '*_1_', seq_suffix])
    mate2_search_string = ''.join([analysis_id, '*_2_', seq_suffix])
    # glob to get the sequence file matches
    mate1_fastq_files = glob.glob(os.path.join(sequence_path, mate1_search_string))
    mate2_fastq_files = glob.glob(os.path.join(sequence_path, mate2_search_string))
    if len(mate1_fastq_files) != 1 and len(mate2_fastq_files) != 1:
        logging.error('Error finding paired-end sequences for flowcell %s lane %s' % (flowcell, lane))
        return
    fastq_files = [mate1_fastq_files[0], mate2_fastq_files[0]]
    # regexp for getting beginning of filename to use as prefix
    # in search for bowtie quals file
    prefix_re = re.compile(r'(.*)' + seq_suffix)    
    quals_param = None    
    for fastq_file in fastq_files:
        m = prefix_re.match(os.path.basename(fastq_file))
        if not m:
            logging.critical("regexp match failed on file %s" % fastq_file)
            continue
        # get the quals param from the bowtie file
        fastq_prefix = m.group(1)
        quals_file = os.path.join(sequence_path, fastq_prefix + bowtie_quals_suffix)            
        if not os.path.exists(quals_file):
            logging.error("quality score file %s not found... skipping %s" %
                          (quals_file, fastq_file))
            return
        
        quals_fhd = open(quals_file)
        tmp_quals_param = quals_fhd.readline()
        quals_fhd.close()
        if quals_param == None:
            quals_param = tmp_quals_param
        elif tmp_quals_param != quals_param:
            logging.error('quality param %s does not match mate pair sequence param %s' % (tmp_quals_param, quals_param))
            return
    # return pair of sequence files and quals param    
    yield fastq_files, quals_param     


def run_tophat(flowcell, lane, fastq_files, quals_param, insert_size, results_path,
               processors=2,
               prefix_fmt="%s_%s"):
    analysis_id = prefix_fmt % (flowcell, lane)
    tophat_cwd = os.path.join(results_path, analysis_id)
    # by default do not overwrite existing results
    if os.path.exists(tophat_cwd):
        logging.debug("tophat results path %s exists, skipping..." % tophat_cwd)
        return
    # skip alignments that already exist
    logging.debug("Preparing to run Tophat on files: %s quals: %s insert_size: %d" % (fastq_files, quals_param, insert_size))
    # make directory for results
    logging.debug('Creating tophat working path %s' % tophat_cwd)
    os.makedirs(tophat_cwd)        
    # options
    tophat_options = {'-p': processors,
                      '--output-dir': tophat_cwd,
                      '--' + quals_param: None,
                      '--mate-inner-dist': insert_size,
                      '--mate-std-dev': 20}                      
    # create the app controller
    tophat_app = TopHat(options=tophat_options)
    # run
    tophat_app.run_paired_end(fastq_files[0], fastq_files[1],
                              output_dir=tophat_cwd, 
                              genome="hg18")
    logging.debug("finished %s" % (analysis_id))


def run_batch_tophat(analyses, sequence_path, results_path):
    for analysis in analyses:
        flowcell, lane, insert_size = analysis
        for res in find_pe_sequence_files(flowcell, lane, sequence_path):
            fastq_files, quals_param = res
            run_tophat(flowcell, lane, fastq_files, quals_param, insert_size, results_path)

def get_all_analyses():
    # query sampleDB for paired-end RNA-seq flowcell/lanes
    analyses = get_sampledb().get_analyses_w_insert_size({'app_type': 'Transcriptome',
                                                          'analysis': 'eland_pair'})
    return analyses

def get_analyses_from_samples(sample_xmlfile):
    analyses = []
    for sample_group in parse_samplegroups_xml(options.samples):
        for sample in sample_group.samples:
            insert_size = get_sampledb().params[sample]['insert_size']
            for fclane in get_sampledb().params[sample]['lanes']:
                logging.debug('    sample: %s lanes: %s' %
                              (sample, fclane))
                analyses.append((fclane[0], fclane[1], insert_size))
    return analyses

if __name__ == '__main__':
    logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s',
                        level=logging.DEBUG)

    optionparser = OptionParser("usage: %prog [options] [-s <samples.xml>] <sequence_path> <results_path>")
    optionparser.add_option("-s", "--samples", dest="samples",
                            help="process only the samples defined in the sample XML file")                            
    (options, args) = optionparser.parse_args()
    if options.samples is not None:
        analyses = get_analyses_from_samples(options.samples)
    else:
        analyses = get_all_analyses()
    # positional arguments
    sequence_path = args[0]
    results_path = args[1]

    # log the command-line arguments
    logging.debug('sequence_path: %s', sequence_path)
    logging.debug('results_path: %s', results_path)
    logging.debug('analyses: %s', analyses)

    # make directory for results
    if not os.path.exists(results_path):
        logging.debug('Creating results path %s' % results_path)
        os.makedirs(results_path)

    # run tophat
    run_batch_tophat(analyses, sequence_path, results_path)